import torch
from torch.utils.data import Dataset
import os
from torchvision import transforms
from PIL import Image


data_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])


class GunDateSet(Dataset):
    def __init__(self, root_dir, sequence_len=17, transform=data_transform):
        self.root_dir = root_dir
        self.sequence_len = sequence_len
        self.transform = transform
        self.sequence = []
        self.label = []
        class_map = {'unfire': 0, 'fire': 1}

        for class_name, label in class_map.items():
            class_dir = os.path.join(root_dir, class_name)
            if not os.path.isdir(class_dir):
                continue
            for seq_folder in os.listdir(class_dir):
                seq_path = os.path.join(class_dir, seq_folder)
                if os.path.isdir(seq_path):
                    self.sequence.append(seq_path)
                    self.label.append(label)

    def __len__(self):
        return len(self.sequence)

    def __getitem__(self, item):
        seq_path = self.sequence[item]
        label = self.label[item]
        frames = []
        frame_files = sorted(os.listdir(seq_path))
        for frame_file in frame_files:
            frame_path = os.path.join(seq_path, frame_file)
            image = Image.open(frame_path).convert('RGB')
            frames.append(image)
        if len(frames) > self.sequence_len:
            frames = frames[:self.sequence_len]
        elif len(frames) < self.sequence_len:
            frames.extend([frames[-1] * (self.sequence_len-len(frames))])
        if self.transform:
            frames = [self.transform(frame) for frame in frames]
        frame_tensor = torch.stack(frames)
        return frame_tensor, torch.tensor(label, dtype=torch.float32)


if __name__ == '__main__':
    root_dir = r'G:\TimeSpeceDate\train'
    dataset = GunDateSet(root_dir=root_dir)
    sample, label = dataset[0]
    print(sample.shape)
    print(label)



